import pandas as pd
from datetime import datetime
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import matplotlib.lines as mlines
import seaborn as sns
import numpy as np
from string import punctuation
import scipy.stats as st


# Load data
df = pd.read_csv('../data/temp/partisan_narratives.csv')
df['narrative'] = df.narrative.apply(lambda x: x.strip())
df['above_zero'] = df.log_or > 0

# Plot
_fig, _ax = plt.subplots(figsize=(20, 8))

_ax.xaxis.label.set_fontsize(25)
_ax.xaxis.labelpad = 20
_ax.set_xlim([-5.0, 5.5])
_ax.xaxis.set_ticks(np.arange(-5.0, 5.5, 1))


_ax.yaxis.label.set_fontsize(25)
_ax.set_axisbelow(True)
_ax.axes.get_yaxis().set_visible(False)
_ax.axes.set_ylim([-5, 5])
_ax.spines['left'].set_visible(False)
_ax.spines['top'].set_visible(False)
_ax.spines['bottom'].set_visible(False)
_ax.spines['right'].set_visible(False)

markers = {True: "o", False: "X"}

_g = sns.scatterplot(data=df,
                     x= 'log_or',
                     y=[0]*20,
                     s=300,
                     style='above_zero',
                     markers=markers,
                     hue='above_zero',
                     palette=['blue', 'red'],
                     alpha=0.8,
                     zorder=2,
                     ax=_ax)

# Point Labels
# The amount that the labels should be placed according to the actual points (x, y)
label_positions_deviations = {
                                'budget end medicare': (-1.8, -3),
                                'american lose unemployment (benefits)': (-3.0, 2.5),
                                'budget cut medicaid': (-0.9, -3.5),
                                'american exhaust unemployment (benefits)': (-1.5, 3.5),
                                'worker exhaust money': (-0.0, -2.7),
                                'budget cut education' : (-0.6, 1.5),
                                'person earn men women': (-0.8, -1.5),
                                'republican hold prisoner': (-0.5, 2.5),
                                'person make dollar': (0.6, -1.9),
                                'men women earn dollar': (0.8, 0.8),

                                # Positive Log_or
                                'money keep money': (-2.5, -2),
                                'government control healthcare': (-2, 2),
                                'doctor practice offense': (-1.0, -3.5),
                                'family keep money': (-0.8, 3.2),
                                'american keep money': (-0.5, -2.5),
                                'person make global warming': (-0.6, 2.5),
                                'god bless men women': (-0.5, -1.2),
                                'american rely oil': (-0.3, 1.5),
                                'government approve healthcare': (-0.7, -3.1),
                                'energy affect family': (-0.1, 3),
                           }

for i in range(df.shape[0]):
    narrative_text = df.narrative[i]
    _ax.annotate(s=df.narrative[i],
                    xy=(df.log_or[i], 0),
                    xytext=(df.log_or[i]+label_positions_deviations[narrative_text][0],
                            label_positions_deviations[narrative_text][1]),
                    arrowprops = dict(
                                        arrowstyle="simple",
                                        shrinkA=5,
                                        shrinkB=15,
                                        color='grey',
                                        alpha=0.2,
#                                      connectionstyle="angle3,angleA=0,angleB=-90"
                                      ),
                    color='black',
                    fontsize = 16,
#                    bbox=dict(pad=5,
#                              facecolor='blue',
#                              alpha=0.2)
                )



plt.xticks(fontsize=18)
plt.yticks(fontsize=18)
_ax.set_axisbelow(True)
_ax.xaxis.grid(color='gray', alpha=0.3, linestyle='dashed')
_g.set(xlabel=None)
_ax.get_legend().remove()

# Output file
plt.savefig('../figures/Figure_3_c.pdf', bbox_inches="tight")
